import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QHBoxLayout, QPushButton, QFileDialog, QComboBox, QLabel, QGroupBox
from PyQt5.QtGui import QPixmap, QImage
import os
from Function import *


class ImageProcessingApp(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()
        self.image = None
        self.processed_image = None
        # 获取当前脚本所在的目录
        self.current_dir = os.path.dirname(os.path.abspath(__file__))

    def initUI(self):
        # 创建主布局
        main_layout = QVBoxLayout()

        # 创建按钮和下拉框布局
        button_layout = QHBoxLayout()

        # 创建按钮
        self.load_button = QPushButton('加载图像', self)
        self.process_button = QPushButton('处理图像', self)
        self.save_button = QPushButton('保存处理后的图像', self)
        self.delete_button = QPushButton('删除处理后的图像', self)
        self.load_button.clicked.connect(self.load_image)
        self.process_button.clicked.connect(self.process_image)
        self.save_button.clicked.connect(self.save_image)
        self.delete_button.clicked.connect(self.delete_processed_image)
        self.save_button.setEnabled(False)
        self.delete_button.setEnabled(False)

        button_layout.addWidget(self.load_button)
        button_layout.addWidget(self.process_button)
        button_layout.addWidget(self.save_button)
        button_layout.addWidget(self.delete_button)

        # 创建下拉框
        self.algorithm_combobox = QComboBox(self)
        self.algorithm_combobox.addItems([
            '线性变换', '对数变换', '幂律变换', '分段线性变换', '直方图均衡化', '直方图规定化',
            '均值滤波', '高斯滤波', '中值滤波', '梯度算子', '拉普拉斯算子', '高通滤波',
            'Sobel边缘检测', 'Prewitt边缘检测', '傅里叶变换', '离散余弦变换',
            '理想低通滤波', '巴特沃斯低通滤波', '理想高通滤波', '拉普拉斯频域滤波',
            '带通滤波', '带阻滤波', '逆滤波', '维纳滤波', '约束最小二乘滤波',
            '迭代盲复原', '凸集投影', '腐蚀', '膨胀', '开运算', '闭运算',
            '顶帽变换', '黑帽变换', '形态学梯度', '骨架提取', '连通域分析',
            '平移', '旋转', '近邻插值缩放', '双线性插值缩放', '镜像', '错切',
            '透视变换', '仿射变换', '校正几何畸变', 'SIFT图像配准', 'ORB图像配准'
        ])
        button_layout.addWidget(self.algorithm_combobox)

        # 创建图像显示布局
        image_layout = QHBoxLayout()

        # 创建原始图像显示组
        original_group = QGroupBox('原始图像')
        original_layout = QVBoxLayout()
        self.original_image_label = QLabel(self)
        original_layout.addWidget(self.original_image_label)
        original_group.setLayout(original_layout)

        # 创建处理后图像显示组
        processed_group = QGroupBox('处理后图像')
        processed_layout = QVBoxLayout()
        self.processed_image_label = QLabel(self)
        processed_layout.addWidget(self.processed_image_label)
        processed_group.setLayout(processed_layout)

        image_layout.addWidget(original_group)
        image_layout.addWidget(processed_group)

        main_layout.addLayout(button_layout)
        main_layout.addLayout(image_layout)

        self.setLayout(main_layout)
        self.setWindowTitle('图像处理应用')
        self.setGeometry(300, 300, 1000, 600)
        self.show()

    def load_image(self):
        file_dialog = QFileDialog()
        file_path, _ = file_dialog.getOpenFileName(self, '选择图像文件', '', '图像文件 (*.png *.jpg *.bmp)')
        if file_path:
            self.image = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
            if self.image is not None:
                self.display_image(self.original_image_label, self.image)
                self.processed_image = None
                self.processed_image_label.clear()
                self.save_button.setEnabled(False)
                self.delete_button.setEnabled(False)

    def process_image(self):
        if self.image is None:
            return
        algorithm = self.algorithm_combobox.currentText()
        result = None
        if algorithm == '线性变换':
            result = linear_transformation(self.image, a=1.5, b=30)
        elif algorithm == '对数变换':
            result = log_transformation(self.image, c=1)
        elif algorithm == '幂律变换':
            result = gamma_transformation(self.image, gamma=0.5)
        elif algorithm == '分段线性变换':
            result = piecewise_linear_transformation(self.image, r1=50, s1=20, r2=200, s2=230)
        elif algorithm == '直方图均衡化':
            result = histogram_equalization(self.image)
        elif algorithm == '直方图规定化':
            reference_path = os.path.join(self.current_dir, 'girl.bmp')  # 使用os.path.join构建路径
            reference = cv2.imread(reference_path, cv2.IMREAD_GRAYSCALE)
            if reference is None:
                print(f"无法读取文件: {reference_path}，请检查文件路径和完整性。")
            else:
                result = histogram_matching(self.image, reference)
        elif algorithm == '均值滤波':
            result = mean_filter(self.image, kernel_size=3)
        elif algorithm == '高斯滤波':
            result = gaussian_filter(self.image, kernel_size=3, sigma=0)
        elif algorithm == '中值滤波':
            result = median_filter(self.image, kernel_size=3)
        elif algorithm == '梯度算子':
            result = gradient_operator(self.image)
        elif algorithm == '拉普拉斯算子':
            result = laplacian_operator(self.image)
        elif algorithm == '高通滤波':
            result = high_pass_filter(self.image)
        elif algorithm == 'Sobel边缘检测':
            result = sobel_edge_detection(self.image)
        elif algorithm == 'Prewitt边缘检测':
            result = prewitt_edge_detection(self.image)
        elif algorithm == '傅里叶变换':
            result = fft_transform(self.image)
        elif algorithm == '离散余弦变换':
            result = dct_transform(self.image)
        elif algorithm == '理想低通滤波':
            result = ideal_low_pass_filter(self.image, cutoff=30)
        elif algorithm == '巴特沃斯低通滤波':
            result = butterworth_low_pass_filter(self.image, cutoff=30, order=2)
        elif algorithm == '理想高通滤波':
            result = ideal_high_pass_filter(self.image, cutoff=30)
        elif algorithm == '拉普拉斯频域滤波':
            result = laplacian_frequency_filter(self.image)
        elif algorithm == '带通滤波':
            result = band_pass_filter(self.image, low_cutoff=10, high_cutoff=30)
        elif algorithm == '带阻滤波':
            result = band_stop_filter(self.image, low_cutoff=10, high_cutoff=30)
        elif algorithm == '逆滤波':
            kernel = np.ones((3, 3), np.float32) / 9
            result = inverse_filter(self.image, kernel, noise_power=0)
        elif algorithm == '维纳滤波':
            kernel = np.ones((3, 3), np.float32) / 9
            result = wiener_filter(self.image, kernel, K=0.01)
        elif algorithm == '约束最小二乘滤波':
            kernel = np.ones((3, 3), np.float32) / 9
            result = constrained_least_squares_filter(self.image, kernel, gamma=0.01)
        elif algorithm == '迭代盲复原':
            result = iterative_blind_deconvolution(self.image, kernel_size=3, iterations=3)
        elif algorithm == '凸集投影':
            kernel = np.ones((3, 3), np.float32) / 9
            result = projection_onto_convex_sets(self.image, kernel, iterations=10)
        elif algorithm == '腐蚀':
            result = erosion(self.image, kernel_size=3)
        elif algorithm == '膨胀':
            result = dilation(self.image, kernel_size=3)
        elif algorithm == '开运算':
            result = opening(self.image, kernel_size=3)
        elif algorithm == '闭运算':
            result = closing(self.image, kernel_size=3)
        elif algorithm == '顶帽变换':
            result = top_hat(self.image, kernel_size=3)
        elif algorithm == '黑帽变换':
            result = black_hat(self.image, kernel_size=3)
        elif algorithm == '形态学梯度':
            result = morphological_gradient(self.image, kernel_size=3)
        elif algorithm == '骨架提取':
            result = skeleton_extraction(self.image)
        elif algorithm == '连通域分析':
            num_labels, labels, stats, centroids = connected_components_analysis(self.image)
            result = np.uint8(labels * (255 / num_labels))
        elif algorithm == '平移':
            result = translation(self.image, x=50, y=50)
        elif algorithm == '旋转':
            result = rotation(self.image, angle=45, scale=1.0)
        elif algorithm == '近邻插值缩放':
            result = nearest_neighbor_scaling(self.image, scale_x=1.5, scale_y=1.5)
        elif algorithm == '双线性插值缩放':
            result = bilinear_scaling(self.image, scale_x=1.5, scale_y=1.5)
        elif algorithm == '镜像':
            result = mirror(self.image, axis=1)
        elif algorithm == '错切':
            result = shear(self.image, shear_factor_x=0.2, shear_factor_y=0.2)
        elif algorithm == '透视变换':
            src_points = np.float32([[0, 0], [self.image.shape[1], 0], [0, self.image.shape[0]], [self.image.shape[1], self.image.shape[0]]])
            dst_points = np.float32([[50, 50], [self.image.shape[1] - 50, 50], [50, self.image.shape[0] - 50], [self.image.shape[1] - 50, self.image.shape[0] - 50]])
            result = perspective_transformation(self.image, src_points, dst_points)
        elif algorithm == '仿射变换':
            src_points = np.float32([[0, 0], [self.image.shape[1], 0], [0, self.image.shape[0]]])
            dst_points = np.float32([[50, 50], [self.image.shape[1] - 50, 50], [50, self.image.shape[0] - 50]])
            result = affine_transformation(self.image, src_points, dst_points)
        elif algorithm == '校正几何畸变':
            src_points = np.float32([[0, 0], [self.image.shape[1], 0], [0, self.image.shape[0]], [self.image.shape[1], self.image.shape[0]]])
            dst_points = np.float32([[50, 50], [self.image.shape[1] - 50, 50], [50, self.image.shape[0] - 50], [self.image.shape[1] - 50, self.image.shape[0] - 50]])
            result = geometric_distortion_correction(self.image, src_points, dst_points, degree=2)
        elif algorithm == 'SIFT图像配准':
            reference_path = os.path.join(self.current_dir, 'girl.bmp')  # 使用os.path.join构建路径
            reference = cv2.imread(reference_path, cv2.IMREAD_GRAYSCALE)
            if reference is None:
                print(f"无法读取文件: {reference_path}，请检查文件路径和完整性。")
            else:
                result = sift_image_registration(self.image, reference)
        elif algorithm == 'ORB图像配准':
            reference_path = os.path.join(self.current_dir, 'girl.bmp')  # 使用os.path.join构建路径
            reference = cv2.imread(reference_path, cv2.IMREAD_GRAYSCALE)
            if reference is None:
                print(f"无法读取文件: {reference_path}，请检查文件路径和完整性。")
            else:
                result = orb_image_registration(self.image, reference)

        if result is not None:
            self.processed_image = result
            self.display_image(self.processed_image_label, result)
            self.save_button.setEnabled(True)
            self.delete_button.setEnabled(True)

    def save_image(self):
        if self.processed_image is not None:
            file_dialog = QFileDialog()
            file_path, _ = file_dialog.getSaveFileName(self, '保存图像', '', '图像文件 (*.png *.jpg *.bmp)')
            if file_path:
                cv2.imwrite(file_path, self.processed_image)

    def delete_processed_image(self):
        self.processed_image = None
        self.processed_image_label.clear()
        self.save_button.setEnabled(False)
        self.delete_button.setEnabled(False)

    def display_image(self, label, image):
        if len(image.shape) == 2:  # 灰度图像
            height, width = image.shape
            qimage = QImage(image.data, width, height, QImage.Format_Grayscale8)
        else:  # 彩色图像
            height, width, channel = image.shape
            bytes_per_line = 3 * width
            qimage = QImage(image.data, width, height, bytes_per_line, QImage.Format_RGB888).rgbSwapped()
        pixmap = QPixmap.fromImage(qimage)
        label.setPixmap(pixmap.scaled(label.width(), label.height()))


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = ImageProcessingApp()
    sys.exit(app.exec_())